import gradio as gr
import pickle
import numpy as np
from PIL import Image
import cv2
from sklearn.datasets import load_digits
import pinecone
from pinecone import Pinecone, ServerlessSpec

# Pinecone配置 - 请替换为你的API密钥
PINECONE_API_KEY = "pcsk_Njyyj_Jx3sovkra1ygNdaXR6vXicgz8s5fiWnS6dbCxTjjr1bwqvZtYeBXxobphF2TiX8"  # 从Pinecone控制台获取
PINECONE_ENV = "us-east-1"  # 例如：us-east1-gcp
INDEX_NAME = "digits-knn"
TOP_K = 5  # 与optimal_knn.py中找到的最佳k值保持一致

# 1. 初始化Pinecone并准备索引
try:
    # 初始化Pinecone客户端
    pc = Pinecone(api_key=PINECONE_API_KEY)
    
    # 检查索引是否存在，不存在则创建
    if INDEX_NAME not in pc.list_indexes().names():
        pc.create_index(
            name=INDEX_NAME,
            dimension=64,  # 8x8图像展平后的维度
            metric="euclidean",  # KNN常用欧氏距离
            spec=ServerlessSpec(
                cloud="aws",  # 根据你的环境修改
                region=PINECONE_ENV
            )
        )
    
    # 连接到索引
    index = pc.Index(INDEX_NAME)
    
    # 加载数据集并导入到Pinecone（仅首次运行时需要）
    digits = load_digits()
    X = digits.data
    y = digits.target
    
    # 检查索引是否已有数据，没有则导入
    if index.describe_index_stats()["total_vector_count"] == 0:
        print("正在将训练数据导入Pinecone...")
        # 批量导入数据（每次100条）
        batch_size = 100
        for i in range(0, len(X), batch_size):
            batch_end = min(i + batch_size, len(X))
            vectors = [
                {
                    "id": f"vec_{j}",
                    "values": X[j].tolist(),
                    "metadata": {"label": int(y[j])}  # 存储标签用于后续投票
                }
                for j in range(i, batch_end)
            ]
            index.upsert(vectors=vectors)
        print(f"成功导入 {len(X)} 条训练数据到Pinecone")
    
    # 验证Pinecone连接（测试样本搜索）
    test_samples = X[0:5]
    test_labels = y[0:5]
    print("\n=== Pinecone验证结果 ===")
    for i in range(5):
        results = index.query(
            vector=test_samples[i].tolist(),
            top_k=TOP_K,
            include_metadata=True
        )
        # 投票获取预测结果
        labels = [int(match["metadata"]["label"]) for match in results["matches"]]
        pred = max(set(labels), key=labels.count)
        print(f"样本{i+1}：真实标签{test_labels[i]} → 预测{pred}")

except Exception as e:
    print(f"Pinecone初始化错误: {str(e)}")
    exit()

# 2. 基于Pinecone的预测函数
def predict_digit(input_data):
    if input_data is None:
        return "请先在画板上绘制一个数字（用白色粗笔）"
    
    # 提取图像数据
    if isinstance(input_data, dict):
        input_data = input_data.get('composite', input_data.get('layers', [None])[-1])
    if input_data is None:
        return "未检测到绘制内容，请重新绘制"
    
    # 处理图像格式
    if not isinstance(input_data, np.ndarray):
        try:
            input_data = np.array(input_data)
        except:
            return "输入格式错误，请重新绘制"
    # 处理RGBA转RGB
    if input_data.shape[2] == 4:
        input_data = input_data[:, :, :3]
    if len(input_data.shape) != 3 or input_data.shape[2] != 3:
        return f"图像格式错误（需3通道RGB），当前: {input_data.shape}"
    
    # 核心预处理（对齐训练数据）
    try:
        # 灰度化 + 低阈值二值化（捕捉淡色笔画）
        gray_img = Image.fromarray(input_data.astype('uint8'), 'RGB').convert('L')
        img_array = np.array(gray_img)
        _, binary_img = cv2.threshold(img_array, 80, 255, cv2.THRESH_BINARY)
        
        # 缩放为8x8 + 颜色反转 + 像素范围对齐
        resized_img = Image.fromarray(binary_img).resize((8, 8), Image.LANCZOS)
        resized_array = np.array(resized_img)
        processed_array = (255.0 - resized_array) * (16 / 255.0)  # 匹配训练数据的0-16像素范围
        processed_array[processed_array < 1] = 0  # 清除背景噪点
        
        # 转换为向量并在Pinecone中搜索
        query_vector = processed_array.flatten().tolist()
        results = index.query(
            vector=query_vector,
            top_k=TOP_K,
            include_metadata=True
        )
        
        # 基于近邻标签投票决定结果
        if not results["matches"]:
            return "未找到相似样本，请重新绘制"
        
        labels = [int(match["metadata"]["label"]) for match in results["matches"]]
        prediction = max(set(labels), key=labels.count)  # 多数投票
        return f"预测结果: {str(prediction)} (基于{TOP_K}个近邻)"
    
    except Exception as e:
        return f"处理出错: {str(e)}"

# 3. 启动服务
with gr.Blocks(title="手写数字识别（Pinecone版）") as iface:
    gr.Markdown("# 手写数字识别（Pinecone云推理）")
    gr.Markdown("绘制要求：1. 用白色粗笔 2. 占满画板 3. 无多余线条")
    
    with gr.Row():
        sketchpad = gr.Sketchpad(label="绘制区域", height=350)
        output = gr.Label(label="识别结果", show_label=True)
    
    with gr.Row():
        submit_btn = gr.Button("提交", variant="primary")
        clear_btn = gr.Button("清除")
    
    # 绑定事件
    submit_btn.click(
        fn=predict_digit,
        inputs=sketchpad,
        outputs=output
    )
    
    clear_btn.click(
        fn=lambda: None,
        inputs=None,
        outputs=sketchpad
    )

if __name__ == "__main__":
    print("正在启动服务... 访问地址：http://127.0.0.1:7864")
    iface.launch(share=False, server_port=7864, server_name="0.0.0.0")